package org.ayachinene.leetcode.utils;

import org.ayachinene.leetcode.structure.TreeNode;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.util.CollectionUtils;

import java.util.List;

public class TreeUtils {

    private static final String TREE_NODE_NULL = "null";

    public static @Nullable TreeNode getTree(@NotNull String s) {
        List<String> elementStrings = CommonUtils.extract(s);

        if (CollectionUtils.isEmpty(elementStrings)) {
            return null;
        }

        List<TreeNode> nodes = elementStrings.stream().map(elementString -> {
            if (elementString.equals(TREE_NODE_NULL)) {
                return null;
            }
            return new TreeNode(Integer.parseInt(elementString));
        }).toList();

        int cntOfNull = 0;
        for (int i = 0; i < nodes.size(); i++) {
            TreeNode node = nodes.get(i);

            if (node == null) {
                cntOfNull++;
                continue;
            }

            int leftIndex = 2 * i + 1 - cntOfNull * 2;
            int rightIndex = 2 * i + 2 - cntOfNull * 2;
            if (leftIndex < nodes.size()) {
                node.left = nodes.get(leftIndex);
            }
            if (rightIndex < nodes.size()) {
                node.right = nodes.get(rightIndex);
            }
        }

        return nodes.get(0);
    }
}
